#include <stdio.h>
#include <algorithm>
#include <fstream>
#include <string>
#include <vector>
#include <iostream>
#include <time.h>

#include "ncnn_tools.h"
#include "benchmark.h"
#include "darknet2ncnn.h"

#define INPUT_SIZE_IMAGE 150
#define PATH_BIN  "/media/thanglmb/DATA/myProject/AI_DL/darknet2ncnn/cameraDemo/models/yolov3-tiny.bin"
#define PATH_PARAM "/media/thanglmb/DATA/myProject/AI_DL/darknet2ncnn/cameraDemo/models/yolov3-tiny.param"
#define PAHT_LABELS "/media/thanglmb/DATA/myProject/AI_DL/darknet2ncnn/cameraDemo/models/coco.names"
struct Object
{
  cv::Rect_<float> rect;
  int label;
  float prob;
};

static void draw_objects(const cv::Mat &bgr, const std::vector<Object> &objects,
                         std::vector<std::string> &labels, cv::Mat &img_out)
{

  cv::Mat image = bgr.clone();

  for (size_t i = 0; i < objects.size(); i++)
  {
    const Object &obj = objects[i];

    fprintf(stderr, "%-3d[%-16s] = %.5f at %.2f %.2f %.2f x %.2f\n", obj.label,
            labels[obj.label].c_str(), obj.prob, obj.rect.x, obj.rect.y,
            obj.rect.width, obj.rect.height);

    cv::rectangle(image, obj.rect, cv::Scalar(255, 0, 0));

    char text[256];
    //sprintf(text, "%s %.1f%%", labels[obj.label].c_str(), obj.prob * 100);
	sprintf(text, "%s ", labels[obj.label].c_str());
    int baseLine = 0;
    cv::Size label_size =
        cv::getTextSize(text, 5, 0, 5, &baseLine);

    int x = obj.rect.x;
    int y = obj.rect.y - label_size.height - baseLine;
    if (y < 0)
      y = 0;
    if (x + label_size.width > image.cols)
      x = image.cols - label_size.width;

    cv::rectangle(
        image,
        cv::Rect(cv::Point(x, y),
                 cv::Size(label_size.width, label_size.height + baseLine)),
        cv::Scalar(255, 255, 255), cv::FILLED);

    cv::putText(image, text, cv::Point(x, y + label_size.height),
                5, 3, cv::Scalar(0, 0, 255));
  }
  //cv::imwrite("test.png", image);
  img_out= image;
}

static int detect_yolov2( cv::Mat &bgr, std::vector<Object> &objects, std::vector<std::string> &labels)
{
  CustomizedNet yolo;
  register_darknet_layer(yolo);
  yolo.load_param(PATH_PARAM);
  yolo.load_model(PATH_BIN);

  //ncnn::Input *input = (ncnn::Input *)yolo.get_layer_from_name("data");
  ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, INPUT_SIZE_IMAGE,INPUT_SIZE_IMAGE);
  const float norm_vals[3] = {1 / 255.0, 1 / 255.0, 1 / 255.0};
  in.substract_mean_normalize(0, norm_vals);

  ncnn::Extractor ex = yolo.create_extractor();
  ex.set_num_threads(2);
  ex.input("data", in);

  ncnn::Mat out;
  ncnn::Blob *out_blob = yolo.get_last_layer_output_blob();
  int result = ex.extract(out_blob->name.c_str(), out);
  if (result != 0)
  {
    printf("ncnn error: %d\n", result);
    return result;
  }

  int img_h = bgr.rows;
  int img_w = bgr.cols;
  
  for (int i = 0; i < out.h; i++)
  {
    const float *values = out.row(i);

    Object object;
    object.label = values[0];
    object.prob = values[1];
    object.rect.x = values[2] * img_w;
    object.rect.y = values[3] * img_h;
    object.rect.width = values[4] * img_w - object.rect.x;
    object.rect.height = values[5] * img_h - object.rect.y;

    objects.push_back(object);
  }

  
  return 0;
}

int main(int argc, char **argv)
{
    std::vector<Object> objects;
    cv::Mat frame;
    
    cv::VideoCapture cap;
    cap.open(0);
    if (!cap.isOpened()) {
        std::cout << "ERROR! Unable to open camera\n";
        return -1;
    }
    while(1)
    {

        cap >> frame;
        if (frame.empty()) {
            std::cout << "ERROR! blank frame grabbed\n";
            break;
        }
        objects.clear();
        std::vector<std::string> labels;
        load_classifier_labels(labels,(char*)PAHT_LABELS );
        double start = static_cast<double>(cv::getTickCount());
        detect_yolov2( frame, objects,labels);
        double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
        std::cout << "spent: " << time << "s " << std::endl;

        cv::Mat out;
        draw_objects(frame, objects, labels,out );
        cv::Mat result;
        cv::resize(out, result, cv::Size(640,480),1);
        cv::imshow("result", result);
        if(cv::waitKey(30)==27) break;
    }
    return 0;
}